/* introspect the VM running in CPU host. We know GPA, and request VMI Daemon to get PA and its data. */
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/mman.h>
#include <stdio.h>
#include <inttypes.h>
#include <getopt.h>

#include <libvmi/libvmi.h>

#define VM_NAME               "vm"
#define DEFAULT_LOG_FILENAME    "agent.log"

int parse_args(int argc, char **argv, uint64_t *vmid, uint64_t *gpa, uint64_t *size, FILE **log_fp) {
    char *output_filename = NULL;

    if ( argc < 2 ) {
        printf("Usage: %s\n", argv[0]);
        printf("\t -o/--output-file <log file>\n");
        printf("\t -v/--vm-id <VM id>\n");
        return 0;
    }

    if ( argc > 2 ) {
        const struct option long_opts[] = {
            {"vm-id", required_argument, NULL, 'v'},
            {"output-file", optional_argument, NULL, 'o'},
            {NULL, 0, NULL, 0}
        };
        const char* opts = "v:g:o:s:";
        int c;
        int long_index = 0;

        while ((c = getopt_long (argc, argv, opts, long_opts, &long_index)) != -1) {
            switch (c) {
                case 'v':
                    *vmid = strtoull(optarg, NULL, 0);
                    break;
                case 'o':
                    output_filename = optarg;
                    break;
                default:
                    printf("Unknown option\n");
                    return -1;
            }
        }
    }

    if (output_filename != NULL) {
        *log_fp = fopen(output_filename, "w+");
    } else {
        *log_fp = fopen(DEFAULT_LOG_FILENAME, "w+");
    }
    if (*log_fp == NULL) {
        printf("cannot write to output file: %s\n", output_filename ? output_filename : DEFAULT_LOG_FILENAME);
        return -1;
    }
    return 0;
}

int test_process_list(vmi_instance_t vmi, FILE *log_fp) {
    vmi_mode_t mode;
    addr_t current_process = 0;
    char *procname = NULL;
    vmi_pid_t pid = 0;
    unsigned long tasks_offset = 0, pid_offset = 0, name_offset = 0;
    addr_t list_head = 0, cur_list_entry = 0, next_list_entry = 0;

    fprintf(log_fp, ">>> luowu. os is linux. try to vmi_get_offset\n");
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_tasks", &tasks_offset) )
        return -1;
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_name", &name_offset) )
        return -1;
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_pid", &pid_offset) )
        return -1;

    /* demonstrate name and id accessors */
    if (VMI_FAILURE == vmi_get_access_mode(vmi, NULL, 0, NULL, &mode))
        return -1;

    if ( VMI_FILE != mode ) {
        uint64_t id = vmi_get_vmid(vmi);
        fprintf(log_fp, "Process listing for VM %s (id=%"PRIu64")\n", VM_NAME, id);
    } else {
        fprintf(log_fp, "Process listing for file %s\n", VM_NAME);
    }
    fflush(log_fp);

    /* get the head of the list */
    if ( VMI_FAILURE == vmi_translate_ksym2v(vmi, "init_task", &list_head) )
        return -1;

    list_head += tasks_offset;
    cur_list_entry = list_head;
    if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry, 0, &next_list_entry)) {
        fprintf(log_fp, "Failed to read next pointer at %"PRIx64"\n", cur_list_entry);
        return -1;
    }

    /* walk the task list */
    while (1) {
        current_process = cur_list_entry - tasks_offset;
        vmi_read_32_va(vmi, current_process + pid_offset, 0, (uint32_t*)&pid);

        procname = vmi_read_str_va(vmi, current_process + name_offset, 0);
        if (!procname) {
            fprintf(log_fp, "Failed to find procname\n");
            return -1;
        }

        /* print out the process name */
        fprintf(log_fp, "[%5d] %s (struct addr:%"PRIx64")\n", pid, procname, current_process);
        fflush(log_fp);
        if (procname) {
            free(procname);
            procname = NULL;
        }

        /* follow the next pointer */
        cur_list_entry = next_list_entry;
        if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry, 0, &next_list_entry)) {
            fprintf(log_fp, "Failed to read next pointer in loop at %"PRIx64"\n", cur_list_entry);
            return -1;
        }

        if (cur_list_entry == list_head)
            break;
    };
    return 0;
}

int main (int argc, char **argv) {
    vmi_instance_t vmi = {0};
    uint64_t gpa = 0, size = 0, vmid = 0;
    FILE *log_fp = NULL;

    int rc = -1;
    rc = parse_args(argc, argv, &vmid, &gpa, &size, &log_fp);
    if (rc != 0) {
        fprintf(log_fp, "Failed to parse args.\n");
        return -1;
    }

    /* initialize the libvmi library */
    fprintf(log_fp, ">>> luowu. try to vmi_init_complete: find kpgd, kaslr, init_task...\n");
    fflush(log_fp);
    if (VMI_FAILURE == vmi_init_complete(&vmi, VM_NAME, VMI_INIT_DOMAINNAME, NULL, VMI_CONFIG_GLOBAL_FILE_ENTRY, NULL, NULL)) {
        fprintf(log_fp, "Failed to init LibVMI library.\n");
        goto error_exit;
    }

    /* init the offset values */
    if (VMI_OS_LINUX != vmi_get_ostype(vmi)) {
        fprintf(log_fp, "Error: ostype is not Linux.\n");
        goto error_exit;
    }

    fprintf(log_fp, ">>> luowu. test... show the %s's process list...\n", VM_NAME);
    fflush(log_fp);
    rc = test_process_list(vmi, log_fp);
    if (rc != 0) {
        fprintf(log_fp, "Failed to get process list.\n");
        return -1;
    }

    fprintf(log_fp, ">>> luowu. DONE...\n");
    fflush(log_fp);

error_exit:
    /* cleanup any memory associated with the LibVMI instance */
    vmi_destroy(vmi);
    return rc;
}
